"""
OA智能体 - 使用智谱AI的ChatGLM实现的OA系统智能助手
可以帮助用户查询审批单据、考勤信息等
支持使用国内开源模型(智谱AI的ChatGLM)进行意图识别
"""
import os
import json
import requests
from typing import Dict, List, Any, Optional
from datetime import datetime

# 配置
OA_API_BASE_URL = "http://localhost:8000"
ZHIPUAI_API_KEY = "855f0ae254b94990abd8a5a978da7823.Ft43vzMrVUcB8PuL"  # 请替换为您的智谱API密钥

# 设置环境变量
os.environ["ZHIPUAI_API_KEY"] = ZHIPUAI_API_KEY

# 导入智谱AI SDK
try:
    import zhipuai
except ImportError:
    print("正在安装zhipuai SDK...")
    os.system("pip install zhipuai>=1.0.7")
    import zhipuai

# OA API客户端
class OAClient:
    def __init__(self, base_url: str):
        self.base_url = base_url
        self.token = None
    
    def login(self, username: str, password: str) -> bool:
        """登录OA系统获取token"""
        try:
            response = requests.post(
                f"{self.base_url}/token",
                json={"username": username, "password": password}
            )
            if response.status_code == 200:
                data = response.json()
                self.token = data["access_token"]
                return True
            return False
        except Exception as e:
            print(f"登录失败: {str(e)}")
            return False
    
    def get_headers(self) -> Dict[str, str]:
        """获取带有认证信息的请求头"""
        if not self.token:
            raise ValueError("未登录，请先调用login方法")
        return {"Authorization": f"Bearer {self.token}"}
    
    def get_user_info(self) -> Dict[str, Any]:
        """获取当前用户信息"""
        response = requests.get(
            f"{self.base_url}/users/me",
            headers=self.get_headers()
        )
        return response.json()
    
    def get_pending_approvals(self) -> Dict[str, Any]:
        """获取待审批的单据"""
        response = requests.get(
            f"{self.base_url}/approvals/pending",
            headers=self.get_headers()
        )
        return response.json()
    
    def get_my_approvals(self) -> Dict[str, Any]:
        """获取我提交的单据"""
        response = requests.get(
            f"{self.base_url}/approvals/my",
            headers=self.get_headers()
        )
        return response.json()
    
    def get_my_attendance(self, year: int, month: int) -> Dict[str, Any]:
        """获取我的考勤记录"""
        response = requests.get(
            f"{self.base_url}/attendance/my",
            params={"year": year, "month": month},
            headers=self.get_headers()
        )
        return response.json()

# 工具函数定义
def get_pending_approvals(oa_client: OAClient) -> str:
    """获取需要我审批的单据"""
    try:
        result = oa_client.get_pending_approvals()
        if result["total"] == 0:
            return "您当前没有需要审批的单据。"
        
        approvals = result["items"]
        response = f"您有 {len(approvals)} 个待审批单据:\n\n"
        
        for i, approval in enumerate(approvals, 1):
            response += f"{i}. 单号: {approval['id']}\n"
            response += f"   标题: {approval['title']}\n"
            response += f"   类型: {approval['type']}\n"
            response += f"   申请人: {approval['applicant']}\n"
            response += f"   提交时间: {approval['submit_time']}\n"
            response += f"   紧急程度: {approval['urgency']}\n"
            if approval['amount']:
                response += f"   金额: ¥{approval['amount']}\n"
            response += f"   内容: {approval['content']}\n\n"
        
        return response
    except Exception as e:
        return f"获取待审批单据时出错: {str(e)}"

def get_my_approvals(oa_client: OAClient) -> str:
    """获取我提交的单据"""
    try:
        result = oa_client.get_my_approvals()
        if result["total"] == 0:
            return "您当前没有提交的单据。"
        
        approvals = result["items"]
        response = f"您提交的单据共 {len(approvals)} 个:\n\n"
        
        for i, approval in enumerate(approvals, 1):
            response += f"{i}. 单号: {approval['id']}\n"
            response += f"   标题: {approval['title']}\n"
            response += f"   类型: {approval['type']}\n"
            response += f"   状态: {approval['status']}\n"
            response += f"   提交时间: {approval['submit_time']}\n"
            if approval['amount']:
                response += f"   金额: ¥{approval['amount']}\n"
            response += f"   内容: {approval['content']}\n\n"
        
        return response
    except Exception as e:
        return f"获取我的单据时出错: {str(e)}"

def get_attendance_info(oa_client: OAClient, year: Optional[int] = None, month: Optional[int] = None) -> str:
    """获取我的考勤信息"""
    try:
        # 如果未指定年月，使用当前年月
        if year is None or month is None:
            now = datetime.now()
            year = now.year
            month = now.month
        
        result = oa_client.get_my_attendance(year, month)
        
        stats = result["statistics"]
        response = f"{year}年{month}月考勤统计:\n"
        response += f"- 正常出勤: {stats['normal']}天\n"
        response += f"- 迟到: {stats['late']}天\n"
        response += f"- 早退: {stats['early_leave']}天\n"
        response += f"- 缺勤: {stats['absent']}天\n"
        response += f"- 请假: {stats['leave']}天\n"
        
        # 检查是否有异常考勤
        if stats['late'] > 0 or stats['early_leave'] > 0 or stats['absent'] > 0:
            response += "\n异常考勤记录:\n"
            for record in result["records"]:
                if record["status"] in ["迟到", "早退", "缺勤"]:
                    response += f"- {record['date']}: {record['status']}"
                    if record["check_in"]:
                        response += f", 签到时间: {record['check_in']}"
                    if record["check_out"]:
                        response += f", 签退时间: {record['check_out']}"
                    response += "\n"
        
        return response
    except Exception as e:
        return f"获取考勤信息时出错: {str(e)}"

def get_user_info(oa_client: OAClient) -> str:
    """获取当前用户信息"""
    try:
        user = oa_client.get_user_info()
        response = "您的个人信息:\n"
        response += f"- 用户名: {user['username']}\n"
        response += f"- 姓名: {user['full_name']}\n"
        response += f"- 邮箱: {user['email']}\n"
        response += f"- 部门: {user['department']}\n"
        response += f"- 职位: {user['position']}\n"
        return response
    except Exception as e:
        return f"获取用户信息时出错: {str(e)}"

# 使用智谱AI的ChatGLM进行意图识别的OA智能体
class OAAgent:
    def __init__(self, oa_client: OAClient):
        self.oa_client = oa_client
        self.commands = {
            "待审批单据": self._get_pending_approvals,
            "我的单据": self._get_my_approvals,
            "考勤信息": self._get_attendance_info,
            "个人信息": self._get_user_info,
            "帮助": self._get_help
        }
        # 初始化智谱AI客户端
        zhipuai.api_key = ZHIPUAI_API_KEY
    
    def _get_pending_approvals(self) -> str:
        return get_pending_approvals(self.oa_client)
    
    def _get_my_approvals(self) -> str:
        return get_my_approvals(self.oa_client)
    
    def _get_attendance_info(self, year=None, month=None) -> str:
        return get_attendance_info(self.oa_client, year, month)
    
    def _get_user_info(self) -> str:
        return get_user_info(self.oa_client)
    
    def _get_help(self) -> str:
        help_text = "OA智能助手可以帮您查询以下信息:\n"
        help_text += "1. 待审批单据 - 查询需要您审批的单据\n"
        help_text += "2. 我的单据 - 查询您提交的单据\n"
        help_text += "3. 考勤信息 - 查询您的考勤记录\n"
        help_text += "4. 个人信息 - 查询您的个人信息\n"
        help_text += "5. 帮助 - 显示此帮助信息\n"
        help_text += "\n您可以直接输入命令或用自然语言提问，例如'我想查看我的考勤'"
        return help_text
    
    def _identify_intent_with_llm(self, query: str) -> Dict[str, Any]:
        """使用智谱AI的ChatGLM模型识别用户意图"""
        try:
            # 导入新版本的智谱AI SDK
            try:
                from zhipuai import ZhipuAI
            except ImportError:
                print("正在安装最新版本的zhipuai SDK...")
                os.system("pip install -U zhipuai")
                from zhipuai import ZhipuAI
            
            # 初始化客户端
            client = ZhipuAI(api_key=ZHIPUAI_API_KEY)
            
            # 构建系统提示和用户查询
            messages = [
                {"role": "system", "content": """你是一个OA系统的意图识别助手，需要从用户输入中识别出用户的意图。
                
                可能的意图类型有：
                1. 查询待审批单据 - 用户想查看需要他审批的单据
                2. 查询我的单据 - 用户想查看自己提交的单据
                3. 查询考勤信息 - 用户想查看自己的考勤记录
                4. 查询个人信息 - 用户想查看自己的个人信息
                5. 帮助 - 用户需要帮助或使用说明
                6. 未知 - 无法确定用户意图
                
                请以JSON格式返回识别结果，格式如下:
                {
                    "intent": "意图类型",
                    "confidence": 0.95,
                    "parameters": {}
                }
                
                其中意图类型必须是以下之一: "查询待审批单据", "查询我的单据", "查询考勤信息", "查询个人信息", "帮助", "未知"
                confidence是置信度，范围0-1
                parameters是可能的参数，例如查询考勤时可能包含年份和月份"""},
                {"role": "user", "content": query}
            ]
            
            # 调用智谱AI的ChatGLM模型
            response = client.chat.completions.create(
                model="glm-4",
                messages=messages,
                temperature=0.1,
                max_tokens=500
            )
            
            # 解析返回结果
            content = response.choices[0].message.content
            
            # 提取JSON部分
            try:
                # 尝试直接解析整个内容
                intent_data = json.loads(content)
            except:
                # 如果失败，尝试提取内容中的JSON部分
                import re
                json_match = re.search(r'({.*})', content, re.DOTALL)
                if json_match:
                    try:
                        intent_data = json.loads(json_match.group(1))
                    except:
                        # 如果仍然失败，返回默认值
                        intent_data = {"intent": "未知", "confidence": 0.0, "parameters": {}}
                else:
                    intent_data = {"intent": "未知", "confidence": 0.0, "parameters": {}}
            
            return intent_data
        except Exception as e:
            print(f"意图识别出错: {str(e)}")
            return {"intent": "未知", "confidence": 0.0, "parameters": {}}
    
    def _fallback_intent_recognition(self, query: str) -> Dict[str, Any]:
        """备用的关键词匹配意图识别方法"""
        query = query.lower()  # 转为小写以便匹配
        
        # 待审批单据相关匹配
        if any(keyword in query for keyword in ["待审批", "需要我审批", "审批单", "待我审批", "等我审批", "审批什么"]):
            return {"intent": "查询待审批单据", "confidence": 0.8, "parameters": {}}
        
        # 我的单据相关匹配
        elif any(keyword in query for keyword in ["我的单据", "我提交的", "我的申请", "我申请的", "查单据", "查下单据", 
                                              "查询单据", "帮我查单据", "帮我查下单据", "我的审批单"]):
            return {"intent": "查询我的单据", "confidence": 0.8, "parameters": {}}
        
        # 考勤信息相关匹配
        elif any(keyword in query for keyword in ["考勤", "打卡", "出勤", "签到", "上班记录", "考勤记录", 
                                              "查考勤", "查询考勤", "我的考勤"]):
            return {"intent": "查询考勤信息", "confidence": 0.8, "parameters": {}}
        
        # 个人信息相关匹配
        elif any(keyword in query for keyword in ["个人信息", "我的信息", "我的资料", "个人资料", "我是谁", 
                                              "我的详情", "用户信息", "查询我的信息"]):
            return {"intent": "查询个人信息", "confidence": 0.8, "parameters": {}}
        
        # 帮助信息相关匹配
        elif any(keyword in query for keyword in ["帮助", "help", "怎么用", "使用说明", "功能", "能做什么", 
                                              "有什么功能", "如何使用"]):
            return {"intent": "帮助", "confidence": 0.8, "parameters": {}}
        
        # 通用查询意图识别
        elif "查" in query and ("单据" in query or "审批" in query or "申请" in query):
            # 根据上下文判断是查询待审批还是我的单据
            if "我的" in query or "我提交" in query or "我申请" in query:
                return {"intent": "查询我的单据", "confidence": 0.7, "parameters": {}}
            else:
                return {"intent": "查询待审批单据", "confidence": 0.7, "parameters": {}}
        
        # 如果是直接命令
        for cmd in self.commands:
            if query == cmd:
                intent_map = {
                    "待审批单据": "查询待审批单据",
                    "我的单据": "查询我的单据",
                    "考勤信息": "查询考勤信息",
                    "个人信息": "查询个人信息",
                    "帮助": "帮助"
                }
                return {"intent": intent_map.get(cmd, "未知"), "confidence": 1.0, "parameters": {}}
        
        # 默认未知意图
        return {"intent": "未知", "confidence": 0.0, "parameters": {}}
    
    def process_query(self, query: str) -> str:
        """处理用户查询"""
        print("正在分析您的问题...")
        
        # 注释掉智谱AI模型调用部分
        # try:
        #     # 尝试使用智谱AI的ChatGLM模型进行意图识别
        #     intent_data = self._identify_intent_with_llm(query)
        #     print(f"意图识别结果: {intent_data}")
        #     
        #     # 如果置信度低或意图未知，使用备用方法
        #     if intent_data["intent"] == "未知" or intent_data["confidence"] < 0.6:
        #         print("使用备用意图识别方法...")
        #         intent_data = self._fallback_intent_recognition(query)
        #         print(f"备用意图识别结果: {intent_data}")
        # except Exception as e:
        #     print(f"智能意图识别失败，使用备用方法: {str(e)}")
        #     intent_data = self._fallback_intent_recognition(query)
        
        # 直接使用关键词匹配方法
        print("使用关键词匹配方法...")
        intent_data = self._fallback_intent_recognition(query)
        print(f"关键词匹配结果: {intent_data}")
        
        # 根据识别出的意图执行相应操作
        intent = intent_data["intent"]
        parameters = intent_data.get("parameters", {})
        
        if intent == "查询待审批单据":
            return self._get_pending_approvals()
        elif intent == "查询我的单据":
            return self._get_my_approvals()
        elif intent == "查询考勤信息":
            year = parameters.get("year")
            month = parameters.get("month")
            return self._get_attendance_info(year, month)
        elif intent == "查询个人信息":
            return self._get_user_info()
        elif intent == "帮助":
            return self._get_help()
        else:
            # 未知意图
            return "抱歉，我不理解您的问题。请尝试使用更明确的表述，例如'查询我的单据'或'待审批单据'。\n" + self._get_help()

def main():
    # 创建OA客户端
    oa_client = OAClient(OA_API_BASE_URL)
    
    # 登录
    print("正在登录OA系统...")
    if not oa_client.login("wangwu", "password"):
        print("登录失败，请检查OA API服务是否启动")
        return
    
    print("登录成功！")
    
    # 创建OA智能体
    agent = OAAgent(oa_client)
    
    # 显示帮助信息
    print("\n" + agent.process_query("帮助"))
    
    # 交互式查询
    print("\n您可以开始查询了，输入'退出'结束对话")
    while True:
        query = input("\n> ")
        if query.lower() in ["退出", "exit", "quit"]:
            break
        
        response = agent.process_query(query)
        print(response)

if __name__ == "__main__":
    main()